from __future__ import absolute_import

import torch
import os
import os.path as osp
import shutil
import errno


def to_numpy(tensor):
    if torch.is_tensor(tensor):
        return tensor.cpu().numpy()
    elif type(tensor).__module__ != 'numpy':
        raise ValueError("Cannot convert {} to numpy array"
                         .format(type(tensor)))
    return tensor


def to_torch(ndarray):
    if type(ndarray).__module__ == 'numpy':
        return torch.from_numpy(ndarray)
    elif not torch.is_tensor(ndarray):
        raise ValueError("Cannot convert {} to torch tensor"
                         .format(type(ndarray)))
    return ndarray


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def mkdir_if_missing(dirname):
    """Creates dirname if it is missing."""
    if not osp.exists(dirname):
        try:
            os.makedirs(dirname)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise


def save_checkpoint(stateDict, epoch, isBest, ckptDir):
    ckptDir = osp.abspath(ckptDir)
    mkdir_if_missing(ckptDir)
    last_path = osp.join(ckptDir, f'ckpt_{epoch}.pt')
    torch.save(stateDict, last_path)
    if isBest:
        shutil.copy(last_path, osp.join(ckptDir, 'ckpt_best.pt'))
        print('best saved at epoch: {}.'.format(epoch))


def save_best_checkpoint(stateDict, ckptDir):
    ckptDir = osp.abspath(ckptDir)
    mkdir_if_missing(ckptDir)
    best_path = osp.join(ckptDir, f'ckpt_best.pt')
    torch.save(stateDict, best_path)


def load_checkpoint(fpath):
    if osp.isfile(fpath):
        checkpoint = torch.load(fpath)
        print("=> Loaded checkpoint '{}'".format(fpath))
        return checkpoint
    else:
        raise ValueError("=> No checkpoint found at '{}'".format(fpath))


def write_to_logfile(pathLogfile, put_str):
    with open(pathLogfile, 'a') as fp:
        fp.write(put_str)
        fp.write('\n')
